#pragma once
#include <mutex>
#include <cassert>
#include <utility>
#include <iostream>
#include <atomic>

namespace MemoryPool
{
    #define MEMORY_POOL_NUM 64
    #define MAX_SLOT_SIZE 512
    #define SLOT_BASE_SIZE 8

    // 非自由链表的槽没有存储next指针，自由链表的槽存储的是next指针
    struct slot
    {
        std::atomic<slot*> next;
    };

    class memoryPool
    {
    public:
        memoryPool(const int blockSize = 4096) : block_size(blockSize)
        {}

        ~memoryPool()
        {
            slot* cur = first_block;
            while (cur)
            {
                slot* next = cur->next;
                operator delete(reinterpret_cast<void*>(cur));
                cur = next;
            }
        }

        void init(int size)
        {
            assert(size > 0);
            first_block = nullptr;
            cur_slot = nullptr;
            last_slot = nullptr;
            free_list.store(nullptr, std::memory_order_relaxed);
            slot_size = size;
        }

        void* allocate()
        {
            slot* st = popFreeList();
            if (st != nullptr)
            {
                return st;
            }
            slot* tmp = nullptr;
            {
                // 如果剩余空间不够了，则向系统申请新的内存块
                std::lock_guard<std::mutex> lck(mutexForNewMemory);
                if (cur_slot >= last_slot)
                {
                    allocateNewBlock();
                }
                slot* tmp = cur_slot;
                // 新申请的内存块的第一个槽被用了，cur_slot需要指向下一个槽的起始地址
                cur_slot = reinterpret_cast<slot*>(reinterpret_cast<char*>(cur_slot) + slot_size);
            }
            return tmp;
        }

        void deallocate(void* ptr)
        {
            if (ptr)
            {
                pushFreeList(reinterpret_cast<slot*>(ptr));
            }
        } 

        void allocateNewBlock()
        {
            void* tmp = operator new(block_size);
            reinterpret_cast<slot*>(tmp)->next = first_block;
            first_block = reinterpret_cast<slot*>(tmp);

            // 需要加上一个指向下一个内存块的指针的大小
            char* body = reinterpret_cast<char*>(tmp) + sizeof(slot*);
            size_t paddingSize = padPointer(body, slot_size);
            cur_slot = reinterpret_cast<slot*>(body + paddingSize);
            // 最后一个槽是不确定的
            last_slot = reinterpret_cast<slot*>(reinterpret_cast<char*>(first_block) + block_size - slot_size + 1);
            free_list.store(nullptr, std::memory_order_relaxed);
        }

        size_t padPointer(char* ptr, size_t align)
        {
            return (align - (reinterpret_cast<size_t>(ptr) % align)) % align;
        }

        bool pushFreeList(slot* ptr)
        {
            while (true)
            {
                slot* oldHead = free_list.load(std::memory_order_relaxed);
                // 将节点指向原来的头节点
                ptr->next.store(oldHead, std::memory_order_relaxed);
                // 将新节点设为头节点
                if (free_list.compare_exchange_weak(oldHead, ptr, std::memory_order_release, std::memory_order_relaxed))
                {
                    return true;
                }
                // CAS失败则重试
            }
        }

        slot* popFreeList()
        {
            while (true)
            {
                slot* oldHead = free_list.load(std::memory_order_relaxed);
                if (oldHead == nullptr)
                {   
                    return nullptr;
                }
                slot* newHead = oldHead->next;
                if (free_list.compare_exchange_weak(oldHead, newHead, std::memory_order_acquire, std::memory_order_acquire))
                {
                    return oldHead;
                }
                // CAS失败则重试
            }
        }
    private:
        slot* first_block;   // 第一个内存块的起始地址
        slot* cur_slot;      // 未被使用的第一个槽的地址
        slot* last_slot;     // 最后一个内存块的末尾地址
        std::atomic<slot*> free_list; // 自由链表的起始地址，即归还的槽全部放入自由链表中
        int block_size;
        int slot_size;
        std::mutex mutexForFreeList; // 自由链表的锁
        std::mutex mutexForNewMemory; // 像系统申请内存时需用的锁
    };

    class HashBucket
    {
    public:
        static memoryPool& getMemoryPool(const int index)
        {
            static memoryPool memoryPoolList[MEMORY_POOL_NUM];
            return memoryPoolList[index];
        }

        static void initMemoryPool()
        {
            for (int i = 0; i < MEMORY_POOL_NUM; i++)
            {
                getMemoryPool(i).init((i + 1) * SLOT_BASE_SIZE);
            }
        }

        static void* useMemory(int size)
        {
            if (size <= 0)
            {
                return nullptr;
            }
            if (size > MAX_SLOT_SIZE)
            {
                return operator new(size);
            }
            return getMemoryPool((size + 7) / SLOT_BASE_SIZE - 1).allocate();
        }

        static void freeMemory(void* ptr, int size)
        {
            if (ptr == nullptr)
            {
                return;
            }
            if (size > MAX_SLOT_SIZE)
            {
                operator delete(ptr);
                return;
            }
            getMemoryPool((size + 7) / SLOT_BASE_SIZE - 1).deallocate(ptr);
        }
    };

    template <typename T, typename ...Args>
    T* newElement(Args&&... args)
    {
        T* ptr = nullptr;
        if ((ptr = reinterpret_cast<T*>(HashBucket::useMemory(sizeof(T)))) != nullptr)
        {
            new(ptr) T(std::forward<Args>(args)...);
        }
        return ptr;
    }

    template <typename T>
    void deleteElement(T* ptr)
    {
        if (ptr)
        {
            ptr->~T();
            HashBucket::freeMemory(reinterpret_cast<void*>(ptr), sizeof(T));
        }
    }
}